import os , traceback
import time

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['DEVICE_ID'] = '1'

# 配置昇腾相关环境变量
os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'
os.environ['ASCEND_DEVICE_ID'] = '1'
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '1'
os.environ['ASCEND_VISIBLE_DEVICES'] = '1'
os.environ['ASCEND_HOME'] = '/usr/local/Ascend'

# 拼接 PATH 和 LD_LIBRARY_PATH（保留原有环境变量值）
os.environ['PATH'] = f"{os.environ['ASCEND_HOME']}/bin:{os.environ.get('PATH', '')}"
os.environ['LD_LIBRARY_PATH'] = f"{os.environ['ASCEND_HOME']}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"

import torch
import torch_npu

from transformers import AutoProcessor, AutoModelForImageTextToText

def main():
    for k ,v in dict(os.environ).items():
        print(k , v)
    model_path = "/data/model/StepFun"
    devicemap = {'':'npu:0'}
    model = AutoModelForImageTextToText.from_pretrained(model_path, device_map=devicemap)
    processor = AutoProcessor.from_pretrained(model_path)
    time.sleep(1000)

if __name__ == '__main__':
    main()